/*-
 * Copyright (c) 2007-2008 SINA Corporation, All Rights Reserved.
 * Author: Zhang shuo <zhangshuo@staff.sina.com.cn>
 * A memcache protocol implement under Zhu Yan's SASE(csf) framework.
 */

#include <sys/types.h>
#include <sys/socket.h>
#include <stdio.h>
#include <netinet/in.h>
#include <sys/time.h>
#include <sys/ioctl.h>
#include <unistd.h>
#include <string.h>
#include <stdlib.h>
#include <dlfcn.h>
#include <pthread.h>
#include <errno.h>
#include <signal.h>
#include <fcntl.h>
#include <sys/socket.h>
#include <arpa/inet.h>


#include "libprotocol.h"
#include "buffer.h"
#include "log.h"
#include "confparser.h"
#include "queue.h"

#include "memcache.h"


/* a definiens of PROTOCOL_NAME which is required by CSF */
#define PROTOCOL_NAME "memcache"

int protocol_init(char *);


/* global funciton pointer of mempool allcation method */
//typedef void *MEMPOOL_ALLOC(const char *, int, int);
//extern _SMALLOC *mempool_alloc;

/* functions to prevent broken pipe */
void init_signal(void);
void handle_signal(int s);

/* functions of csf (gossip module) */
static void *memcache_session_start(CONN_INFO *);
static int memcache_session_entry(void *, CONN_INFO *, void *, void *, int);
static void memcache_session_end(CONN_INFO *, void *);
static void mg_data_cleaner(void *);
ssize_t mg_request_responder(CONN_INFO *, void *, int);

/* functions of protocol_basic module */
static ssize_t process_header(PROTOCOL *);
void *buf_transfer(char *, size_t);
int p_responser(PROTOCOL *);

static ssize_t writen(int, const void *, size_t);
static int gossip_connect(char *, int);
void register_self(void);

/* functions of mc_gossip command processing */
static char *get_header_field(PROTOCOL *, int);
size_t cmd_set(PROTOCOL *);
size_t cmd_add(PROTOCOL *);
size_t cmd_replace(PROTOCOL *);
size_t cmd_append(PROTOCOL *);
size_t cmd_prepend(PROTOCOL *);
size_t cmd_cas(PROTOCOL *);
size_t cmd_get(PROTOCOL *);
size_t cmd_gets(PROTOCOL *);
size_t cmd_delete(PROTOCOL *);
size_t cmd_incr(PROTOCOL *);
size_t cmd_decr(PROTOCOL *);
size_t cmd_stats(PROTOCOL *);
size_t cmd_flush_all(PROTOCOL *);
size_t cmd_version(PROTOCOL *);
size_t cmd_quit(PROTOCOL *);
size_t cmd_reg(PROTOCOL *);
size_t cmd_disreg(PROTOCOL *);
size_t cmd_gspset(PROTOCOL *);


/* "command <-> function" array */
static CMD_FUNCTION cmd_func_array[] = {
	{"set", cmd_set},
	{"add", cmd_add},
	{"replace", cmd_replace},
	{"append", cmd_append},
	{"prepend", cmd_prepend},
	{"cas", cmd_cas},
	{"get", cmd_get},
	{"gets", cmd_gets},
	{"delete", cmd_delete},
	{"incr", cmd_incr},
	{"decr", cmd_decr},
	{"stats", cmd_stats},
	{"flush_all", cmd_flush_all},
	{"version", cmd_version},
	{"quit", cmd_quit},
	{"gsp_reg", cmd_reg},
	{"gsp_disreg", cmd_disreg},
    {"gsp_set", cmd_gspset},
	{NULL, NULL}
};


/* config buffer (sl.conf) */
static char init_host[CONF_ITEM_LEN + 1];
static char current_host[CONF_ITEM_LEN + 1];
static int init_port;
static int current_port;
static int default_ttl;
static int peer_id;


static struct conf_int_config conf_int_array[] = {
	{"init_port", &init_port},
	{"current_port", &current_port},
	{"default_ttl", &default_ttl},
	{"peer_id", &peer_id},
	{0, 0}
};

static struct conf_str_config conf_str_array[] = {
	{"init_host", init_host},
	{"current_host", current_host},
	{0, 0}
};

/* =====================================================================
 *         [protocol_basic module dispatching functions]
 * the functions below are dispatched by protocol_basic module 
 * under a specific order.
 * =================================================================== */

/* 
 * process the formatted header. send the command to the function pointed
 * by the pointer which is matched. this function is automatically called
 * by protocol_basic module when header fully read.
 */
static ssize_t 
process_header(PROTOCOL *handler)
{
	CMD_FUNCTION *cmd_function;
	ssize_t func_res = -1;
	int str_len;
	
	
	if (get_header_count(handler) <= 0)return 0;
	
	/* get the command array address */
	cmd_function = cmd_func_array;
	while (cmd_function->command != NULL)
	{
		if (get_header_fields(handler)[0] == NULL)
		{
			str_len = 0;
		}
		else
		{
			str_len = strlen(get_header_fields(handler)[0]);
		}
		
		/* compare the input command with the function array's items */
		if (strncmp(get_header_fields(handler)[0], cmd_function->command, str_len) == 0)
		{
			/* note that this is a function pointer. type: CMD_FUNCTION*  */
			func_res = cmd_function->func_ptr(handler);
			break;
		}
		cmd_function++;
	}
	
	return func_res;
}


/* 
 * to copy the buffer to another place. this function is automatically called
 * when buffer is full.
 */
void *
buf_transfer(char *data, size_t len)
{
	CSF_UNUSED_ARG(data);
	CSF_UNUSED_ARG(len);
	
	//printf("&&&transfer: %ld bytes transfered, data is: \n%s\n", len, data);
	
	return NULL;
}


/* 
 * when all the data received, this function is called to response the request
 */
int 
p_responser(PROTOCOL *handler)
{
	int r;
	
    r = submit_request(0, get_custom_p1(handler), handler, 
								mg_data_cleaner, mg_request_responder, REQUEST_NONE);
    if (r < 0) 
	{
		WLOG_ERR("pipeline error. %d\n", r);
		return P_RESPONSER_ERR;
	}

    /* set the submitting success flag */
	((MC_INFO *)get_user_data(handler))->is_submit = 1;

	return P_RESPONSER_OK;
}


/* =====================================================================
 *                 [CSF dispatching functions]
 * the functions below are dispatched by CSF protocol module
 * =================================================================== */


static void *
memcache_session_start(CONN_INFO *cip)
{
	PROTOCOL *new_handler;
	MC_INFO *user_data;

    #ifdef MC_GOSSIP_DEBUG
    printf("==========================================\n");
    #endif
	
	/* create a new protocol_basic instance */
	new_handler = pb_init("\r\n", " ", 
						buf_transfer, p_responser, process_header);
						
	/* create user_data to store anlyzed infomation */
	user_data = (MC_INFO *)smalloc(sizeof(MC_INFO));		
	if (user_data == NULL || new_handler == NULL)
    {
        WLOG_ERR("not enough memory.");
        return NULL;
    }
			
	/* save this instance's pointer to the main handler */
    new_handler->is_submit = 0;
	set_user_data(new_handler, user_data);
						
	/* save the cip structure to the main handler */
	set_custom_p2(new_handler, cip);
	
    return new_handler;
}


/* get and process TCP socket data */
static int 
memcache_session_entry(void *rqstpp, CONN_INFO *cip, void *cli_info, void *data, int len)
{
	PROTOCOL *p_handler;
	int p_entry_res;

	CSF_UNUSED_ARG(cip);
	
	/* get the handler created in memcache_session_start */
	p_handler = (PROTOCOL *)cli_info;
	
	if (p_handler == NULL || len <= 0)
	{
        return PROTOCOL_DISCONNECT;
	}
	
	set_custom_p1(p_handler, rqstpp);
	
	/* if the buffer is FULL, try unitl it responses CONTINUE or ERROR */
	do
	{	p_entry_res = protocol_entry(p_handler, data, (size_t)len);
		
	}while(p_entry_res == LEN_BUF_AGAIN);
	
	if (p_entry_res == ERROR)
	{
		return PROTOCOL_DISCONNECT;
	}
		
    return 0;
}


static void 
memcache_session_end(CONN_INFO *cip, void *ptr)
{
    /* no use */
    PROTOCOL *handler;
	handler = (PROTOCOL *)ptr;
	
	CSF_UNUSED_ARG(cip);

    if (((MC_INFO *)get_user_data(handler)) == NULL)
    {
        WLOG_ERR("ptr is NULL when to free.");
        return;
    }
	
	if (((MC_INFO *)get_user_data(handler))->is_submit != 1)
	{
		
		sfree((MC_INFO *)get_user_data(handler));
		protocol_clean(handler);
	}
	
    #ifdef MC_GOSSIP_DEBUG
	printf("protocol end!!\n");
    #endif

    return;
}


/* =====================================================================
 *                 [CSF dispatching functions]
 * the functions below are dispatched by CSF pipeline module (MOD).
 * =================================================================== */


static void
mg_data_cleaner(void *ptr)
{
    PROTOCOL *handler;
	handler = (PROTOCOL *)ptr;

    if (ptr == NULL)
    {
        WLOG_ERR("ptr is NULL when to free.");
        return;
    }
	
    sfree((MC_INFO *)get_user_data(handler));
	protocol_clean(handler);
	
	#ifdef MC_GOSSIP_DEBUG
    printf("gossip clean!!\n");
    #endif
    
    return;
}

ssize_t 
mg_request_responder(CONN_INFO *cip, void *ptr, int flag)
{
    PROTOCOL *handler;
	handler = (PROTOCOL *)ptr;
    
    CSF_UNUSED_ARG(cip);
	CSF_UNUSED_ARG(flag);
    
    #ifdef MC_GOSSIP_DEBUG
    printf("gossip response!!");
    #endif
    
    return (RESPONDER_OK | RESPONDER_DISCONNECT);
}


/* =====================================================================
 *                 [CSF dispatching functions]
 * the functions below are for CSF function resgister.
 * =================================================================== */


/* automatically call by csf */
int 
protocol_init(char *proto_name)
{
    int r;
    pthread_t tid;

    set_protocol_session_start(memcache_session_start);
	set_protocol_session_entry(memcache_session_entry);
	set_protocol_session_end(memcache_session_end);

    /* load string config */
	r = load_conf(NULL, proto_name, NULL, conf_str_array);
    
    /* load integer config */
    r = load_conf(NULL, proto_name, conf_int_array, NULL);

	if (r != 0){
	    WLOG_ERR(" load_conf() failed!");
	}

    init_signal();

    /* use a thread to delay the register */
    if (pthread_create(&tid, NULL, (void *)register_self, NULL) != 0){
		WLOG_ERR("creating thread failed.");
		return PROTOCOL_DISCONNECT;
	}
    
    return PROTOCOL_OK;
}


void register_self()
{
    int fd;
    uint64_t msg_id;
    char cmd_buf[MAX_CMD_LEN + 1];
    
    FILE *file;
    uint64_t rand_num;

    file = fopen("/dev/random", "r");
    fread(&rand_num, sizeof(uint64_t), 1, file);
    fclose(file);
    
    usleep(2000000);
    
    /* register myself to init_host:init_port */
    if (strncmp(init_host, current_host, strlen(init_host)) != 0
        || init_port != current_port)
    {
        #ifdef MC_GOSSIP_DEBUG
        printf("\n\n\n\n(msg)[%d] registering to: %s:%d\n", peer_id, init_host, init_port);
        #endif
        
        fd = gossip_connect(init_host, init_port);
        msg_id = rand_num; 

        snprintf(cmd_buf, MAX_CMD_LEN, "gsp_reg %s %d %d %d %lu %d %d\r\n", 
                current_host, current_port, peer_id,
                default_ttl, msg_id, peer_id, peer_id);  
        
        #ifdef MC_GOSSIP_DEBUG
        printf("\t%s\n", cmd_buf);
        #endif

        writen(fd, cmd_buf, strlen(cmd_buf));

        close(fd);
    }

    return;
}


/* 
 * capture signal 13, broken pipe 
 * the default behavior is ignore the signal
 */
void 
init_signal(void)
{
    signal(SIGPIPE,handle_signal);
}


/* do nothing with the signal and reset the signal option */
void 
handle_signal(int s) 
{ 
    CSF_UNUSED_ARG(s);
    init_signal();
}






/* =====================================================================
 *                [mc_gossip command processing functions]
 * each command refer to a specified function.
 * =================================================================== */







/* ------------------------------------
 *            storage command 
 * -----------------------------------*/

size_t 
cmd_set(PROTOCOL *handler)
{
	MC_INFO *user_data;
	
	/* copy the data to specific position of the struct */
	user_data = get_user_data(handler);
	
	user_data->cmd = _cmd_set;
	user_data->key_count = 1;		/* because of SET can only has ONE key */
	user_data->keys[0] = get_header_field(handler, 1);
	user_data->flag = get_header_field(handler, 2);
	user_data->exptime = get_header_field(handler, 3);
	
	if (get_header_field(handler, 4) != NULL)
	{
		user_data->bytes = (size_t)safe_atoll(get_header_field(handler, 4));
	}
	else
	{
		user_data->bytes = 0;
	}
	
	/* be always not forget to return the entity data length */
	return user_data->bytes;
}

size_t 
cmd_add(PROTOCOL *handler)
{
	size_t res;
	MC_INFO *user_data;
	
	/* parsing process is similar to parsing cmd_set */
	res = cmd_set(handler);
	
	/* but the command name is different, we correct it */
	user_data = get_user_data(handler);
	user_data->cmd = _cmd_add;
	
	return res;
}

size_t 
cmd_replace(PROTOCOL *handler)
{
	size_t res;
	MC_INFO *user_data;
	
	/* parsing process is similar to parsing cmd_set */
	res = cmd_set(handler);
	
	/* but the command name is different, we correct it */
	user_data = get_user_data(handler);
	user_data->cmd = _cmd_replace;
	
	return res;
}

size_t 
cmd_append(PROTOCOL *handler)
{
	size_t res;
	MC_INFO *user_data;
	
	/* parsing process is similar to parsing cmd_set */
	res = cmd_set(handler);
	
	/* but the command name is different, we correct it */
	user_data = get_user_data(handler);
	user_data->cmd = _cmd_append;
	
	return res;
}

size_t 
cmd_prepend(PROTOCOL *handler)
{
	size_t res;
	MC_INFO *user_data;
	
	/* parsing process is similar to parsing cmd_set */
	res = cmd_set(handler);
	
	/* but the command name is different, we correct it */
	user_data = get_user_data(handler);
	user_data->cmd = _cmd_prepend;
	
	return res;
}

size_t 
cmd_cas(PROTOCOL *handler)
{
	size_t res;
	MC_INFO *user_data;
	
	/* parsing process is similar to parsing cmd_set */
	res = cmd_set(handler);
	
	/* but the command name is different, we correct it */
	user_data = get_user_data(handler);
	user_data->cmd = _cmd_cas;
	user_data->cas_unqiue = get_header_field(handler, 5);
	
	return res;
}


/* ------------------------------------
 *          retrieval command 
 * -----------------------------------*/


size_t 
cmd_get(PROTOCOL *handler)
{
	MC_INFO *user_data;
	int i;
	
	/* copy the data to specific position of the struct */
	user_data = get_user_data(handler);
	
	user_data->cmd = _cmd_get;
	user_data->key_count = get_header_count(handler) - 1;
	
	for (i = 0; i < user_data->key_count && i < MAX_KEYS; i++)
	{
		user_data->keys[i] = get_header_field(handler, i + 1);
	}

	/* be always not forget to return the entity data length */
	return 0;
}

size_t 
cmd_gets(PROTOCOL *handler)
{
	MC_INFO *user_data;
	
	/* parsing process is similar to parsing cmd_set */
	cmd_get(handler);
	
	/* but the command name is different, we correct it */
	user_data = get_user_data(handler);
	user_data->cmd = _cmd_gets;
	
	return 0;
}


/* ------------------------------------
 *         deletion command 
 * -----------------------------------*/
 

size_t 
cmd_delete(PROTOCOL *handler)
{
	MC_INFO *user_data;
	
	/* copy the data to specific position of the struct */
	user_data = get_user_data(handler);
	
	user_data->cmd = _cmd_delete;
	user_data->key_count = 1;
	user_data->keys[0] = get_header_field(handler, 1);
	
	if (get_header_field(handler, 2) != NULL)
	{
		user_data->time = get_header_field(handler, 2);
	}
	
	/* be always not forget to return the entity data length */
	return 0;
}


/* ------------------------------------
 *          inplace command 
 * -----------------------------------*/


size_t 
cmd_incr(PROTOCOL *handler)
{
	MC_INFO *user_data;
	
	/* copy the data to specific position of the struct */
	user_data = get_user_data(handler);
	
	user_data->cmd = _cmd_incr;
	user_data->key_count = 1;
	user_data->keys[0] = get_header_field(handler, 1);
	user_data->value = (uint64_t)safe_atoll(get_header_field(handler, 2));

	/* be always not forget to return the entity data length */
	return 0;
}

size_t 
cmd_decr(PROTOCOL *handler)
{
	MC_INFO *user_data;
	
	/* copy the data to specific position of the struct */
	user_data = get_user_data(handler);
	
	user_data->cmd = _cmd_decr;
	user_data->key_count = 1;
	user_data->keys[0] = get_header_field(handler, 1);
	user_data->value = (uint64_t)safe_atoll(get_header_field(handler, 2));

	/* be always not forget to return the entity data length */
	return 0;
}


/* ------------------------------------
 *    gossip reg & disreg
 * -----------------------------------*/

size_t cmd_reg(PROTOCOL *handler)
{
	MC_INFO *user_data;
	
	/* copy the data to specific position of the struct */
	user_data = get_user_data(handler);
	
	user_data->cmd = _cmd_reg;
	user_data->host = get_header_field(handler, 1);
	user_data->port = (int)safe_atol(get_header_field(handler, 2));
    user_data->peer_id = (int)safe_atol(get_header_field(handler, 3));
	user_data->ttl = (int)safe_atol(get_header_field(handler, 4));
    user_data->msg_id = (uint64_t)safe_atoll(get_header_field(handler, 5));
    user_data->sender_id = (int)safe_atol(get_header_field(handler, 6));
    user_data->src_id = (int)safe_atol(get_header_field(handler, 7));
	
	return 0;
}

size_t cmd_disreg(PROTOCOL *handler)
{
	MC_INFO *user_data;
	
	/* copy the data to specific position of the struct */
	user_data = get_user_data(handler);
	
	user_data->cmd = _cmd_disreg;
    user_data->peer_id = (int)safe_atol(get_header_field(handler, 1));
	user_data->ttl = (int)safe_atol(get_header_field(handler, 2));
    user_data->msg_id = (uint64_t)safe_atoll(get_header_field(handler, 3));
    user_data->sender_id = (int)safe_atol(get_header_field(handler, 4));
    user_data->src_id = (int)safe_atol(get_header_field(handler, 5));
	
	return 0;
}

size_t cmd_gspset(PROTOCOL *handler)
{
	MC_INFO *user_data;
	
	/* copy the data to specific position of the struct */
	user_data = get_user_data(handler);
	
	user_data->cmd = _cmd_gspset;
    user_data->key_count = 1;
    user_data->keys[0] = get_header_field(handler, 1);
    user_data->flag = get_header_field(handler, 2);
    user_data->exptime = get_header_field(handler, 3);
    user_data->bytes = (size_t)safe_atoll(get_header_field(handler, 4));
	user_data->ttl = (int)safe_atol(get_header_field(handler, 5));
    user_data->msg_id = (uint64_t)safe_atoll(get_header_field(handler, 6));
    user_data->sender_id = (int)safe_atol(get_header_field(handler, 7));
    user_data->src_id = (int)safe_atol(get_header_field(handler, 8));
	
	return 0;
}


/* ------------------------------------
 *    other command (not support now)
 * -----------------------------------*/


size_t 
cmd_stats(PROTOCOL *handler)
{
	handler = NULL;
	return 0;
}

size_t 
cmd_flush_all(PROTOCOL *handler)
{
	handler = NULL;
	return 0;
}

size_t 
cmd_version(PROTOCOL *handler)
{
	handler = NULL;
	return 0;
}

size_t 
cmd_quit(PROTOCOL *handler)
{
	handler = NULL;
	return 0;
}


/* ------------------------------------
 *    |||||||||||||||||||||||||||||
 * -----------------------------------*/


/* 
 * to prevent from segment fault or other error.
 * here provides a safety method to accessing req_header.fields[n] 
 */
static char *
get_header_field(PROTOCOL *handler, int idx)
{
	if (idx < get_header_count(handler))
	{
		return get_header_fields(handler)[idx];
	}
	else
	{
		return NULL;
	}
}


/*
 * connect to remote server
 */
static int 
gossip_connect(char *host, int port)
{
    int sockfd;
    struct sockaddr_in	servaddr;
    
    sockfd = socket(AF_INET, SOCK_STREAM, 0);

    bzero(&servaddr, sizeof(servaddr));
    servaddr.sin_family      = AF_INET;
    servaddr.sin_addr.s_addr = inet_addr(host);
    servaddr.sin_port        = htons(port);
	    
    if (connect(sockfd, (struct sockaddr *)&servaddr, sizeof(servaddr)) < 0)
    {
		close(sockfd);
        return -1;
    }

    return sockfd;
}


/* 
 * write N bytes data to specific sockfd.
 * used in block-IO.
 */
static ssize_t	
writen(int fd, const void *vptr, size_t n)
{
	size_t		nleft;
    size_t		nwritten;
	const char	*ptr;

	ptr = vptr;
	nleft = n;
	while (nleft > 0) 
    {
		if ((nwritten = write(fd, ptr, nleft)) <= 0) 
        {
			if (errno == EINTR)
				nwritten = 0;		/* and call write() again */
			else
				return(-1);			/* error */
		}

		nleft -= nwritten;
		ptr   += nwritten;
	}
	return(n);
}




